#include "myslam/backend.h"
#include "myslam/algorithm.h"
#include "myslam/feature.h"
#include "myslam/map.h"
#include "myslam/mappoint.h"

namespace myslam {

Backend::Backend() {
    backend_running_.store(true);
    backend_thread_ = std::thread(std::bind(&Backend::BackendLoop, this));
}

void Backend::UpdateMap() {
    std::unique_lock<std::mutex> lock(data_mutex_);
    map_update_.notify_one();
}

void Backend::Stop() {
    backend_running_.store(false);
    map_update_.notify_one();
    backend_thread_.join();
}

//后端并没有进行边缘化，而是对不需要优化的mappoint、frame直接丢弃
void Backend::BackendLoop() {
    while (backend_running_.load()) {
        std::unique_lock<std::mutex> lock(data_mutex_);
        map_update_.wait(lock);

        /// 后端仅优化激活的Frames和Landmarks
        Map::KeyframesType active_kfs = map_->GetActiveKeyFrames();
        Map::LandmarksType active_landmarks = map_->GetActiveMapPoints();
        OptimizeNothing(active_kfs, active_landmarks);
    }
}

void Backend::OptimizeNothing(Map::KeyframesType &keyframes, Map::LandmarksType &landmarks){
    // K 和左右外参
    double fx = cam_left_->fx_;
    double fy = cam_left_->fy_;
    double cx = cam_left_->cx_;
    double cy = cam_left_->cy_;
    SE3 left_ext = cam_left_->pose();
    SE3 right_ext = cam_right_->pose();

    //找到所有要优化的关键帧和路标点，把指向他们的指针保存在一个顺序容器中
    //把关键帧的位姿、地图点的坐标也保存在顺序容器中，避免迭代中反复调用frame->Pose()和mappoint->Pos()，这俩都会加锁
    //再创建一个容器保存上次迭代前的位姿，用于在total_cost增长时回滚位姿估计
    std::vector<Frame::Ptr> kfs;
    std::vector<SE3> kfs_pose;
    std::vector<SE3> kfs_pose_pre;
    std::vector<MapPoint::Ptr> lms;
    std::vector<Vec3> lms_pos;
    std::vector<Vec3> lms_pos_pre;
    for(auto &kf:keyframes) {
        kfs.emplace_back(kf.second);
        kfs_pose.emplace_back(kf.second->Pose());
    }
    for(auto &lm:landmarks) if(!lm.second->is_outlier_) {
        lms.emplace_back(lm.second);
        lms_pos.emplace_back(lm.second->Pos());
    }


    //创建一个顺序容器保存所有参与优化的feature，以及这些feature对应的cost，用于判断outlier
    std::vector<Feature::Ptr> feats;
    std::vector<double> costs;
    double cost_th=5; //cost的阈值
    double total_cost=0;
    double last_total_cost=0;

    int NF = kfs.size()*6;
    int NP = lms.size()*3;

    //H矩阵分块
    Eigen::MatrixXd B;
    B.resize(NF,NF);
    std::vector<Mat33> C(NP/3);
    std::vector<Mat33> C_inv(NP/3);
    Mat33 I = Mat33::Identity(); //单位阵，用于计算C的逆矩阵
    Eigen::MatrixXd E;
    E.resize(NF,NP);
    Eigen::MatrixXd E_t;
    E_t.resize(NP,NF);
    //b分块
    Eigen::VectorXd v;
    v.resize(NF);
    Eigen::VectorXd w;
    w.resize(NP);

    //delta_x分块
    Eigen::VectorXd dxf;
    dxf.resize(NF);
    Eigen::VectorXd dxp;
    dxp.resize(NP);

    //开始迭代
    //结束迭代的情况分两种
    //dx无解或cost增加：结束迭代，不更新pose
    //dx过小或iter=10：更新pose，统计outlier，结束迭代
    int iter=0;
    while(1){
        //在顺序容器中遍历所有的路标点，遍历路标点对应的所有的观测
        //找到这个观测对应的关键帧
        //求该关键点的重投影误差
        //求该观测对对应的路标点和关键帧的J矩阵
        //计算Jt*J和Jt*e，累加到H和b上

        B.setZero();
        for(auto & block:C) block.setZero();
        E.setZero();
        v.setZero();
        w.setZero();
        costs.clear();
        feats.clear();
        total_cost=0;
        for(int landmark_index=0; landmark_index<lms.size();landmark_index++){
            auto obs=lms[landmark_index]->GetObs();
            Vec3 landmark_pos = lms_pos[landmark_index];  ///路标点坐标///
            for(auto &ob:obs){
                //如果feature是空的，或者是外点，或者不关联任何frame，那就跳过它
                if( (ob.lock()==nullptr) || (ob.lock()->is_outlier_) || (ob.lock()->frame_.lock()==nullptr) )
                    continue;
                auto feat=ob.lock();  ///特征点指针///
                auto frame=ob.lock()->frame_.lock(); ///关键帧指针///
                SE3 cam_ext = feat->is_on_left_image_ ? left_ext:right_ext;


                //找到这个frame在顺序容器中的序号
                int frame_index=-1;
                for(int idx=0;idx<kfs.size();idx++)
                    if(frame->id_ == kfs[idx]->id_){
                        frame_index=idx;
                        break;
                    }


                //如果在关键帧列表里找不到这个观测对应的帧，就跳过它
                if(frame_index==-1) continue;

                //计算路标点在对应的相机坐标系（左目或右目）中的坐标
                Vec3 pc;
                pc = cam_ext * kfs_pose[frame_index] * landmark_pos;
                double X = pc(0, 0);
                double Y = pc(1, 0);
                double Z = pc(2, 0);

                //计算重投影误差
                Vec2 e(feat->position_.pt.x - (fx * X / Z + cx), feat->position_.pt.y - (fy * Y / Z + cy));

                //TODO: 实现huber核

                //total_cost用于判断迭代结束
                total_cost += e.squaredNorm();
                //costs用于判断outlier
                costs.emplace_back(e.squaredNorm());
                feats.emplace_back(feat);

                //计算重投影误差对body位姿的Jacobian，其实和对相机位姿的Jacobian一样??
                //注意这里其实是-J
                Mat26 J_body;
                J_body << -fx / Z, 0, fx * X / Z / Z, fx * X * Y / Z / Z, -fx - fx * X * X / Z / Z, fx * Y / Z,
                        0, -fy / Z, fy * Y / Z / Z, fy + fy * Y * Y / Z / Z, -fy * X * Y / Z / Z, -fy * X / Z;
                //计算重投影误差对路标点位置的jacbobian，也是-J
                Mat23 J_landmark;
                J_landmark << J_body.block<2, 3>(0, 0) *
                        cam_ext.rotationMatrix() * kfs_pose[frame_index].rotationMatrix();

                B.block<6,6>(6*frame_index,6*frame_index) += J_body.transpose() * J_body;
                C[landmark_index] += J_landmark.transpose() * J_landmark;
                E.block<6,3>(6*frame_index,3*landmark_index) += J_body.transpose() * J_landmark;

                v.block<6,1>(6*frame_index,0) += -J_body.transpose() * e;
                w.block<3,1>(3*landmark_index,0) += -J_landmark.transpose() * e;
            }
            //完成一个landmark对应的所有feature的遍历
        }
        //完成了对所有landmark的遍历

        //避免H不正定，偷懒的LM
        for(int i=0;i<B.rows();i++)
            B(i,i) += 0.1;
        for(auto &block:C) for(int i=0;i<block.rows();i++)
            block(i,i) += 0.1;

        //利用C的稀疏性计算其逆矩阵
        for(int i=0;i<C.size();i++)
            C_inv[i] = C[i].ldlt().solve(I);

        //利用C_inv的稀疏性计算B-E*C_inv*E.transpose()
        E_t = E.transpose();
        for(int i=0;i<lms.size();i++)
            B = B - E.block(0,3*i,NF,3) * C_inv[i] * E_t.block(3*i,0,3,NF);

        //利用C_inv的稀疏性计算v-E*C_inv*w
        for(int i=0;i<lms.size();i++)
            v = v - E.block(0,3*i,NF,3) * C_inv[i] * w.block<3,1>(3*i,0);
        clock_t tt = clock();
        dxf = B.ldlt().solve(v);

        //利用C_inv稀疏性计算dxp
        w -= E_t*dxf;
        for(int i=0;i<lms.size();i++)
            dxp.block<3,1>(3*i,0) = C_inv[i] * w.block<3,1>(3*i,0);

        std::cout<<"TOTAL COST AT ITERATION "<<iter<<" IS "<<total_cost<<". DX MODULUS IS "<<dxf.norm()+dxp.norm()<<std::endl;

        //如果dx无解或者cost增长，结束迭代，pose回滚
        if( (std::isnan(dxf[0])) || (std::isnan(dxp[0])) || (iter>0&&total_cost>=last_total_cost) ){
            lms_pos = lms_pos_pre;
            kfs_pose = kfs_pose_pre;
            break;
        }

        //更新关键帧pose
        kfs_pose_pre = kfs_pose;
        for(int i=0; i<kfs_pose.size(); i++)
            kfs_pose[i] = SE3::exp(dxf.block(6*i, 0, 6, 1)) * kfs_pose[i];

        //更新路标点坐标
        lms_pos_pre = lms_pos;
        for(int i=0; i<lms_pos.size(); i++)
            lms_pos[i] += dxp.block(3*i,0,3,1);

        //更新cost和iter
        iter++;
        last_total_cost=total_cost;

        //如果iter=10，或者dx已经太小，结束迭代
        if( (iter==10) || (dxf.norm()+dxp.norm()<2.5) ){
            break;
        }
    }

    //调整cost的阈值，至少让inlier的比例达到0.5
    int cnt_outlier = 0, cnt_inlier = 0;
    int iteration = 0;
    while (iteration < 5) {
        cnt_outlier = 0;
        cnt_inlier = 0;
        // determine if we want to adjust the outlier threshold
        for (double cost:costs) {
            if(cost>cost_th) cnt_outlier++;
            else cnt_inlier++;
        }
        double inlier_ratio = cnt_inlier / double(cnt_inlier + cnt_outlier);
        if (inlier_ratio > 0.5) {
            break;
        }
        else{
            cost_th *= 2;
            iteration++;
        }
    }

    //把outlier feature移除掉
    for (int i=0;i<feats.size();i++) {
        if (costs[i]>cost_th) {
            feats[i]->is_outlier_ = true;
            feats[i]->map_point_.lock()->RemoveObservation(feats[i]);
        }
        else feats[i]->is_outlier_ = false;
    }

    LOG(INFO) << "Outlier/Inlier in optimization: " << cnt_outlier << "/" <<cnt_inlier;

    //把kfs_pose、lms_pos里保存的修改传入到kfs、lms
    for (int i=0;i<kfs.size();i++)
        kfs[i]->SetPose(kfs_pose[i]);

    for (int i=0;i<lms.size();i++)
        lms[i]->SetPos(lms_pos[i]);
}

}  // namespace myslam